import os
import uuid
import threading
from datetime import datetime
from typing import List, Dict, Any
import logging
import requests

class MusicGenerator:
    """音乐生成器 - 使用Replicate API避免本地模型问题"""
    
    def __init__(self):
        """初始化音乐生成器"""
        self._lock = threading.Lock()
        self.output_dir = os.path.join(os.path.dirname(__file__), '..', 'static', 'musicgen_outputs')
        self._ensure_output_dir()
        
        # Replicate API配置
        self.model_version = "stereo-large"
        self.output_format = "mp3"
        self.normalization_strategy = "peak"
    
    def _ensure_output_dir(self):
        """确保输出目录存在"""
        os.makedirs(self.output_dir, exist_ok=True)
    
    def generate_music(self, descriptions: List[str], duration: float = 5.0,
                      guidance_scale: float = 3.0, temperature: float = 1.0) -> Dict[str, Any]:
        """使用Replicate API生成音乐"""
        try:
            import replicate
            
            # 验证参数
            if not descriptions or not isinstance(descriptions, list):
                return {
                    "success": False,
                    "error": "请提供有效的音乐描述列表"
                }
            
            if len(descriptions) > 8:  # 限制批量数量
                descriptions = descriptions[:8]
            
            if not (1.0 <= duration <= 30.0):
                return {
                    "success": False,
                    "error": "时长必须在1-30秒之间"
                }
            
            generated_files = []
            
            with self._lock:
                for i, description in enumerate(descriptions):
                    try:
                        # 构建Replicate输入参数
                        inputs = {
                            "prompt": description,
                            "duration": min(duration, 30),
                            "model_version": self.model_version,
                            "output_format": self.output_format,
                            "normalization_strategy": self.normalization_strategy
                        }
                        
                        # 调用Replicate API
                        output = replicate.run(
                            "meta/musicgen:671ac645ce5e552cc63a54a2bbff63fcf798043055d2dac5fc9e36a837eedcfb",
                            input=inputs
                        )
                        
                        if output:
                            # 下载音频文件
                            filename = f"music_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}.mp3"
                            filepath = os.path.join(self.output_dir, filename)
                            
                            # 下载文件
                            audio_data = requests.get(output).content
                            with open(filepath, 'wb') as f:
                                f.write(audio_data)
                            
                            generated_files.append({
                                "description": description,
                                "filename": filename,
                                "filepath": os.path.abspath(filepath),
                                "duration": duration,
                                "sample_rate": 32000,  # Replicate默认采样率
                                "file_size": os.path.getsize(filepath)
                            })
                            
                    except Exception as e:
                        continue  # 跳过失败的请求
                
                if not generated_files:
                    return {
                        "success": False,
                        "error": "所有音乐生成请求失败"
                    }
                
                return {
                    "success": True,
                    "generated_files": generated_files,
                    "total_files": len(generated_files),
                    "model": "Replicate MusicGen API",
                    "sampling_rate": 32000
                }
                
        except ImportError:
            return {
                "success": False,
                "error": "Replicate库未安装，请运行: pip install replicate"
            }
        except Exception as e:
            error_msg = str(e)
            if "replicate" in error_msg.lower():
                return {
                    "success": False,
                    "error": "Replicate API调用失败，请检查网络连接"
                }
            else:
                return {
                    "success": False,
                    "error": f"音乐生成失败: {error_msg}"
                }
    
    def get_supported_params(self) -> Dict[str, Any]:
        """获取支持的参数信息"""
        return {
            "model": "Replicate MusicGen API",
            "sampling_rate": 32000,
            "max_descriptions": 8,
            "duration_range": {
                "min": 1.0,
                "max": 30.0,
                "default": 5.0
            },
            "model_version": {
                "type": "str",
                "description": "模型版本",
                "options": ["melody", "medium", "large", "stereo-large"],
                "default": "stereo-large"
            }
        }
    
    def cleanup_old_files(self, max_files: int = 50) -> Dict[str, Any]:
        """清理旧的生成文件"""
        try:
            if not os.path.exists(self.output_dir):
                return {"success": True, "cleaned": 0}
            
            files = [f for f in os.listdir(self.output_dir) if f.endswith('.wav')]
            files.sort(key=lambda x: os.path.getctime(os.path.join(self.output_dir, x)))
            
            to_remove = max(0, len(files) - max_files)
            cleaned = 0
            
            for i in range(to_remove):
                filepath = os.path.join(self.output_dir, files[i])
                try:
                    os.remove(filepath)
                    cleaned += 1
                except:
                    pass
            
            return {
                "success": True,
                "cleaned": cleaned,
                "remaining": len(files) - cleaned
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }

# 创建全局实例
music_generator = MusicGenerator()
